import os
import sys
import cv2
import numpy as np


# 添加各个模块对应的路径
sys.path.append(os.path.abspath("yolo"))
sys.path.append(os.path.abspath("depth_estimate"))
sys.path.append(os.path.abspath("nlp_model"))
sys.path.append(os.path.abspath("caption_module"))

import yolo.yoloDetector as Yolodetector
import depth_estimate.depth_estimator as DephtEstimator
import nlp_model.qwen as NLPmodel
import caption_module.imageCaption as ImageCaptioner
import pyttsx3

# 参数
TestOnImage = False

# 模型路径
yolo_model_path = "yolo/models/yolov8l.pt"

depth_model_dir = "depth_estimate/models"
depth_model_name = "mono+stereo_640x192"

modeldir = "caption_module/models/"
encoder_path = modeldir + "11encoder_weights.pth"
decoder_path = modeldir + "11decoder_weights.pth"
word_map_path = modeldir + "WORDMAP_coco_5_cap_per_img_5_min_word_freq.json"

# 初始化
# print("Loading models...")
depth_estimator = DephtEstimator.DepthEstimator(depth_model_name, depth_model_dir)
detector = Yolodetector.YoloDetector(yolo_model_path)
image_caption = ImageCaptioner.ImageCaption(encoder_path, decoder_path, word_map_path)


def proccess_frame(img):
    # 推理获取 yolo描述 和 caption 描述
    depthMap = depth_estimator.get_depth(img)
    detections = detector.detect(img, depth_map=depthMap) 
    caption = image_caption.caption(img)

    # 使用 大语言模型 进行语句重组
    discription = NLPmodel.getNlpResponds(detections, caption)

    # 输出
    print(discription)
    pyttsx3.speak(discription)


# 测试
# print("Testing...")
if __name__ == "__main__":
    if TestOnImage:
        img_path = "./depth_estimate/test.png"
        img = cv2.imread(img_path)
        proccess_frame(img)
    else:  # 在视频上进行测试
        cap = cv2.VideoCapture(0)
        while True:
            ret, img = cap.read()
            if not ret:
                break
            proccess_frame(img)
            cv2.imshow("img", img)
            if cv2.waitKey(1) == ord("q"):
                break
        cap.release()
        cv2.destroyAllWindows
